package com.dlkw.easyroom.compiler;

import androidx.room.AutoMigration;
import androidx.room.ColumnInfo;
import androidx.room.Dao;
import androidx.room.Database;
import androidx.room.DatabaseView;
import androidx.room.Entity;

import com.dlkw.easyroom.common.Consts;
import com.dlkw.easyroom.compiler.classkeeper.BaseClass;
import com.dlkw.easyroom.compiler.classkeeper.DaoClass;
import com.dlkw.easyroom.compiler.classkeeper.EntityClass;
import com.dlkw.easyroom.compiler.classkeeper.ViewClass;
import com.squareup.javapoet.AnnotationSpec;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.FieldSpec;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.ParameterizedTypeName;
import com.squareup.javapoet.TypeName;
import com.squareup.javapoet.TypeSpec;
import com.squareup.javapoet.WildcardTypeName;

import java.io.IOException;
import java.lang.annotation.Annotation;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Completion;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.element.Element;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.VariableElement;
import javax.tools.Diagnostic;

public class EasyRoomProcessor extends AbstractProcessor {

    private DataBaseDescriptor dataBaseDescriptor;

    private final Set<DaoClass> daoClasses = new HashSet<>();
    private final Set<EntityClass> entityClasses = new HashSet<>();
    private final Set<ViewClass> databaseViewClasses = new HashSet<>();

    private static final ClassName BASE_DAO_CLASS = ClassName.get("com.dlkw.easyroom", "BaseDao");

    @Override
    public synchronized void init(ProcessingEnvironment processingEnvironment) {
        super.init(processingEnvironment);
        Tools.init(processingEnvironment);
        Tools.log("------------init------------");

        Map<String, String> options = processingEnvironment.getOptions();

        daoClasses.clear();
        entityClasses.clear();
        databaseViewClasses.clear();

        String version = options.get("roomDbVersion");
        String roomDbClass = options.get("roomDbClass");
        String exportSchema = options.get("roomDbExportSchema");

        dataBaseDescriptor = new DataBaseDescriptor(null == roomDbClass ? Consts.DEFAULT_APP_DATABASE_NAME : roomDbClass);
        dataBaseDescriptor.setVersion(null == version ? 1 : Integer.parseInt(version));
        dataBaseDescriptor.setExportSchema(null == exportSchema || Boolean.parseBoolean(exportSchema));
    }

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return new HashSet<>(Arrays.asList(
                Dao.class.getCanonicalName(),
                Entity.class.getCanonicalName(),
                DatabaseView.class.getCanonicalName(),
                AutoMigration.class.getCanonicalName()
        ));
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latest();
    }

    @Override
    public Set<String> getSupportedOptions() {
        Set<String> origin = super.getSupportedOptions();
        if (null == origin) {
            origin = new HashSet<>();
        }
        Set<String> set = new HashSet<>(origin);
        set.add("roomDbClass");
        set.add("roomDbVersion");
        set.add("roomDbExportSchema");
        return set;
    }

    @Override
    public Iterable<? extends Completion> getCompletions(Element element, AnnotationMirror annotationMirror, ExecutableElement executableElement, String s) {
        return super.getCompletions(element, annotationMirror, executableElement, s);
    }

    @Override
    public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnvironment) {
        if (null == set || set.isEmpty()) {
            return false;
        }
        daoClasses.clear();
        entityClasses.clear();
        databaseViewClasses.clear();

        collectDaoElements(roundEnvironment, daoClasses, Dao.class);
        collectEntityElements(roundEnvironment, entityClasses, Entity.class);
        collectViewElements(roundEnvironment, databaseViewClasses, DatabaseView.class);

        dataBaseDescriptor.setDaoClasses(daoClasses);
        dataBaseDescriptor.setEntityClasses(entityClasses);
        dataBaseDescriptor.setViewClasses(databaseViewClasses);

        generateDbClassKeeper();

        ClassName className = ClassName.bestGuess(dataBaseDescriptor.getRoomDbClass());
        String packageName = className.packageName();
        TypeSpec.Builder databaseClassBuilder = buildDatabaseClassBuilder(className.simpleName(), daoClasses);
        databaseClassBuilder.addAnnotation(buildDataBaseAnnotationSpec());

        generateBaseClassKeeper(daoClasses);
        generateBaseClassKeeper(entityClasses);
        generateBaseClassKeeper(databaseViewClasses);

        generateEntityProperties(entityClasses);

        try {
            Tools.writeToFiler(JavaFile.builder(packageName, databaseClassBuilder.build()).build());
            Tools.log("EasyRoom Generate success!");
        } catch (IOException e) {
            Tools.log(Diagnostic.Kind.ERROR, "EasyRoom Generate error:" + e.getMessage());
            throw new RuntimeException(e);
        }
        return false;
    }

    private void generateEntityProperties(Set<EntityClass> entityClasses) {
        ClassName propertyClass = ClassName.bestGuess("com.dlkw.easyroom.Property");
        for (EntityClass entityClass : entityClasses) {
            TypeElement ele = entityClass.getTypeElement();
            ClassName eleClass = ClassName.get(ele);
            TypeSpec.Builder builder = TypeSpec.classBuilder("$" + eleClass.simpleName() + "Property")
                    .addModifiers(Modifier.PUBLIC);

            List<VariableElement> fields = Tools.getFields(ele);
            for (VariableElement field : fields) {
                ColumnInfo columnInfo = field.getAnnotation(ColumnInfo.class);
                String name = null == columnInfo ? field.getSimpleName().toString() : columnInfo.name();
                builder.addField(
                        FieldSpec.builder(
                                        propertyClass,
                                        field.getSimpleName().toString(),
                                        Modifier.PUBLIC,
                                        Modifier.STATIC,
                                        Modifier.FINAL
                                )
                                .initializer("new $T($S, $T.class)", propertyClass, name, field.asType())
                                .build()
                );
            }

            try {
                Tools.writeToFiler(JavaFile.builder(eleClass.packageName(), builder.build()).build());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private void generateBaseClassKeeper(Set<? extends BaseClass> daoClasses) {
        for (BaseClass daoClass : daoClasses) {
            try {
                Tools.writeToFiler(daoClass.toJavaFile());
            } catch (IOException e) {
                Tools.log(Diagnostic.Kind.ERROR, "EasyRoom Generate error:" + e.getMessage());
                throw new RuntimeException(e);
            }
        }
    }

    private void generateDbClassKeeper() {
        ClassName roomDatabaseClassName = ClassName.get("androidx.room", "RoomDatabase");
        WildcardTypeName roomDbSubtype = WildcardTypeName.subtypeOf(roomDatabaseClassName);
        TypeName classOfRoomDb = ParameterizedTypeName.get(
                ClassName.get(Class.class),
                roomDbSubtype
        );

        String dbClass = dataBaseDescriptor.getRoomDbClass();
        ClassName className = ClassName.bestGuess(dbClass);

        TypeSpec.Builder builder = TypeSpec.classBuilder(Consts.DB_CLASS_KEEPER);
        builder.addModifiers(Modifier.PUBLIC, Modifier.FINAL);
        builder.addField(FieldSpec.builder(classOfRoomDb, "DB_CLASS").addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                .initializer("$T.class", className)
                .build());
        JavaFile javaFile = JavaFile.builder(Consts.EASY_ROOM_PACKAGE, builder.build())
                .build();
        try {
            Tools.writeToFiler(javaFile);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private void collectDaoElements(RoundEnvironment roundEnvironment, Set<DaoClass> daoClasses, Class<? extends Annotation> daoClass) {
        Set<? extends Element> daoElements = roundEnvironment.getElementsAnnotatedWith(daoClass);
        if (null != daoElements) {
            for (Element element : daoElements) {
                daoClasses.add(new DaoClass((TypeElement) element));
            }
        }
    }

    private void collectEntityElements(RoundEnvironment roundEnvironment, Set<EntityClass> daoClasses, Class<? extends Annotation> daoClass) {
        Set<? extends Element> daoElements = roundEnvironment.getElementsAnnotatedWith(daoClass);
        if (null != daoElements) {
            for (Element element : daoElements) {
                daoClasses.add(new EntityClass((TypeElement) element));
            }
        }
    }

    private void collectViewElements(RoundEnvironment roundEnvironment, Set<ViewClass> daoClasses, Class<? extends Annotation> daoClass) {
        Set<? extends Element> daoElements = roundEnvironment.getElementsAnnotatedWith(daoClass);
        if (null != daoElements) {
            for (Element element : daoElements) {
                daoClasses.add(new ViewClass((TypeElement) element));
            }
        }
    }

    private TypeSpec.Builder buildDatabaseClassBuilder(String dbClassName, Set<DaoClass> daoClasses) {
        ClassName roomDatabaseClassName = ClassName.get(Tools.getTypeElement("androidx.room.RoomDatabase"));
        TypeSpec.Builder builder = TypeSpec.classBuilder(dbClassName)
                .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
                .superclass(roomDatabaseClassName);
        Map<TypeElement, MethodSpec> methodMap = new HashMap<>();
        Map<String, MethodSpec> daoMethodMap = new HashMap<>();
        for (DaoClass daoClass : daoClasses) {
            TypeElement typeElement = daoClass.getTypeElement();
            String name = typeElement.getSimpleName().toString();
            int i = 2;
            String temp = name;
            while (daoMethodMap.containsKey(temp)) {
                temp = name + i;
                i++;
            }
            MethodSpec methodSpec = MethodSpec.methodBuilder(temp)
                    .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
                    .returns(ClassName.get(typeElement))
                    .build();
            methodMap.put(typeElement, methodSpec);
            daoMethodMap.put(temp, methodSpec);
            daoClass.setDaoGetterName(temp);
        }
        for (DaoClass daoClass : daoClasses) {
            TypeElement typeElement = daoClass.getTypeElement();
            builder.addMethod(methodMap.get(typeElement));
        }
        return builder;
    }

    private static CodeBlock buildEntitiesArrayBlock(Set<EntityClass> elements) {
        if (elements.isEmpty()) {
            return CodeBlock.of("{}");
        }

        return CodeBlock.builder()
                .add("{\n")
                .indent()
                .add(elements.stream()
                        .map(element -> CodeBlock.of("$T.class", ClassName.get(element.getTypeElement())))
                        .collect(CodeBlock.joining(",\n")))
                .unindent()
                .add("\n}")
                .build();
    }

    private static CodeBlock buildViewsArrayBlock(Set<ViewClass> elements) {
        if (elements.isEmpty()) {
            return CodeBlock.of("{}");
        }

        return CodeBlock.builder()
                .add("{\n")
                .indent()
                .add(elements.stream()
                        .map(element -> CodeBlock.of("$T.class", ClassName.get(element.getTypeElement())))
                        .collect(CodeBlock.joining(",\n")))
                .unindent()
                .add("\n}")
                .build();
    }


    private AnnotationSpec buildDataBaseAnnotationSpec() {
        AnnotationSpec.Builder annotationBuilder = AnnotationSpec.builder(ClassName.get(Database.class))
                .addMember("version", "$L", dataBaseDescriptor.getVersion())
                .addMember("exportSchema", "$L", dataBaseDescriptor.isExportSchema());
        if (null != dataBaseDescriptor.getViewClasses() && !dataBaseDescriptor.getViewClasses().isEmpty()) {
            annotationBuilder.addMember("views", buildViewsArrayBlock(dataBaseDescriptor.getViewClasses()));
        }
        if (null != dataBaseDescriptor.getEntityClasses() && !dataBaseDescriptor.getEntityClasses().isEmpty()) {
            annotationBuilder.addMember("entities", buildEntitiesArrayBlock(dataBaseDescriptor.getEntityClasses()));
        }
        return annotationBuilder.build();
    }
}